﻿// SPDX-FileCopyrightText: 2020 Frans van Dorsselaer
//
// SPDX-License-Identifier: GPL-3.0-only

using System.ComponentModel;
using System.Net;
using System.Runtime.InteropServices;
using Microsoft.Win32.SafeHandles;
using Windows.Win32;
using Windows.Win32.Foundation;
using Windows.Win32.Storage.FileSystem;

namespace Usbipd;

sealed partial class DeviceFile : IDisposable
{
    public DeviceFile(string fileName)
    {
        FileHandle = PInvoke.CreateFile(fileName, (uint)(FILE_ACCESS_RIGHTS.FILE_READ_DATA | FILE_ACCESS_RIGHTS.FILE_WRITE_DATA),
            FILE_SHARE_MODE.FILE_SHARE_READ | FILE_SHARE_MODE.FILE_SHARE_WRITE,
            null, FILE_CREATION_DISPOSITION.OPEN_EXISTING, FILE_FLAGS_AND_ATTRIBUTES.FILE_FLAG_OVERLAPPED, null);

        try
        {
            if (FileHandle.IsInvalid)
            {
                throw new Win32Exception("CreateFile");
            }
            if (!PInvoke.SetFileCompletionNotificationModes(FileHandle, (byte)PInvoke.FILE_SKIP_COMPLETION_PORT_ON_SUCCESS))
            {
                throw new Win32Exception("SetFileCompletionNotificationModes");
            }
            BoundHandle = ThreadPoolBoundHandle.BindHandle(FileHandle);
        }
        catch
        {
            FileHandle.Dispose();
            throw;
        }
    }

    readonly SafeFileHandle FileHandle;
    readonly ThreadPoolBoundHandle BoundHandle;

    public HANDLE DangerousGetHandle()
    {
        ObjectDisposedException.ThrowIf(FileHandle.IsClosed, this);
        return (HANDLE)FileHandle.DangerousGetHandle();
    }

    public Task<uint> IoControlAsync(uint ioControlCode, byte[]? input, byte[]? output, bool exactOutput = true)
    {
        var taskCompletionSource = new TaskCompletionSource<uint>(TaskCreationOptions.RunContinuationsAsynchronously);

        unsafe // DevSkim: ignore DS172412
        {
            void OnCompletion(uint errorCode, uint numBytes, NativeOverlapped* nativeOverlapped)
            {
                BoundHandle.FreeNativeOverlapped(nativeOverlapped);
                if ((WIN32_ERROR)errorCode == WIN32_ERROR.ERROR_SUCCESS)
                {
                    if (exactOutput && ((output?.Length ?? 0) != numBytes))
                    {
                        taskCompletionSource.SetException(
                            new ProtocolViolationException($"DeviceIoControl returned {numBytes} bytes, expected {output?.Length ?? 0}"));
                    }
                    else
                    {
                        taskCompletionSource.SetResult(numBytes);
                    }
                }
                else
                {
                    taskCompletionSource.SetException(new Win32Exception((int)errorCode));
                }
            }

            var nativeOverlapped = BoundHandle.AllocateNativeOverlapped(OnCompletion, null, new object?[] { input, output });
            if (PInvoke.DeviceIoControl(FileHandle, ioControlCode, input, output, out var bytesReturned, nativeOverlapped))
            {
                // synchronous completion (will not call OnCompletion due to FILE_SKIP_COMPLETION_PORT_ON_SUCCESS).
                BoundHandle.FreeNativeOverlapped(nativeOverlapped);
                return exactOutput && ((output?.Length ?? 0) != bytesReturned)
                    ? throw new ProtocolViolationException($"DeviceIoControl returned {bytesReturned} bytes, expected {output?.Length ?? 0}")
                    : Task.FromResult(bytesReturned);
            }
            else
            {
                var errorCode = (WIN32_ERROR)Marshal.GetLastPInvokeError();
                if (errorCode == WIN32_ERROR.ERROR_IO_PENDING)
                {
                    // asynchronous pending (will eventually call OnCompletion).
                    return taskCompletionSource.Task;
                }
                else
                {
                    // synchronous error (will not call OnCompletion).
                    BoundHandle.FreeNativeOverlapped(nativeOverlapped);
                    throw new Win32Exception((int)errorCode);
                }
            }
        }
    }

    public Task<uint> IoControlAsync<T>(T ioControlCode, byte[]? input, byte[]? output, bool exactOutput = true) where T : Enum
    {
        return IoControlAsync((uint)(object)ioControlCode, input, output, exactOutput);
    }

    bool IsDisposed;
    public void Dispose()
    {
        if (!IsDisposed)
        {
            BoundHandle.Dispose();
            FileHandle.Dispose();
            IsDisposed = true;
        }
    }
}
